from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import numpy as np
import torch
import re
from transformers import BertTokenizer, pipeline
from model.base_model import BertForIntentClassification
from peft import PeftModel

app = FastAPI(title="电商客服意图识别API")

# 加载模型和分词器
model_dir = "../model/saved_models"
label_classes = np.load("../data/label_classes.npy")

# 文本清洗函数
def clean_text(text):
    text = re.sub(r'[^\u4e00-\u9fff\s]', '', text).strip()
    return text

# 加载分词器
tokenizer = BertTokenizer.from_pretrained(model_dir)

# 加载基础模型
base_model = BertForIntentClassification.from_pretrained(
    "bert-base-chinese",
    num_labels=len(label_classes),
    ignore_mismatched_sizes=True
)

# 加载LoRA适配器
model = PeftModel.from_pretrained(base_model, model_dir)
model.eval()

# 创建推理管道
classifier = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    device=0 if torch.cuda.is_available() else -1,
    return_all_scores=True
)

# 定义请求和响应格式
class IntentRequest(BaseModel):
    text: str

class IntentResponse(BaseModel):
    intent: str
    confidence: float
    top_intents: list

# 健康检查接口
@app.get("/health")
async def health_check():
    return {"status": "healthy"}

# 推理接口
@app.post("/predict", response_model=IntentResponse)
async def predict_intent(request: IntentRequest):
    try:
        # 文本预处理
        cleaned_text = clean_text(request.text)
        
        # 模型推理
        results = classifier(cleaned_text)[0]
        
        # 处理结果
        sorted_results = sorted(results, key=lambda x: x['score'], reverse=True)
        top_intent = sorted_results[0]
        
        return {
            "intent": label_classes[int(top_intent['label'].split('_')[-1])],
            "confidence": float(top_intent['score']),
            "top_intents": [
                {
                    "intent": label_classes[int(res['label'].split('_')[-1])],
                    "confidence": float(res['score'])
                } for res in sorted_results[:3]
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))